//
//  BinaryRedBlackTree.swift
//  BinaryTree
//
//  Created by 郭红旗 on 2021/12/6.
//

import UIKit


class BinaryRedBlackTreeManager{
    typealias comparatorClosure = (_ e1:Int,_ e2:Int) -> Int
    private var comparator:comparatorClosure
    /// 根节点
    var root:BinaryRedBlackTree?
    /// 比较器
    init(comparator:@escaping comparatorClosure) {
        self.comparator = comparator
    }
    
    func createNodeOne(){
        let one = BinaryRedBlackTreeManager.createNode(num: 70)
        let two = BinaryRedBlackTreeManager.createNode(num: 80)
        let three = BinaryRedBlackTreeManager.createNode(num: 76)
        let four = BinaryRedBlackTreeManager.createNode(num: 88)
        let five = BinaryRedBlackTreeManager.createNode(num: 60)
        let six = BinaryRedBlackTreeManager.createNode(num: 50)
        let seven = BinaryRedBlackTreeManager.createNode(num: 62)
        
        one.right = two
        two.parent = one
        
        one.left = five
        five.parent = one
        
        two.right = four
        four.parent = two
        
        two.left = three
        three.parent = two
        
        five.left = six
        six.parent = five
        
        five.right = seven
        seven.parent = five
        
        root = one

    }
    
    static func createNode(num:Int,isBlack:Bool = true) -> BinaryRedBlackTree {
        let one = BinaryRedBlackTree.init(num: num, parentNode: nil)
        isBlack ? one.blackNode() : one.redNode()
        return one
        
    }
    
    /// 添加节点
    func add(value:Int) {
        guard let tree = root else {
            root = BinaryRedBlackTree.init(num: value, parentNode: nil)
            root!.blackNode()
            afterAdd(node: root!)
            return
        }
        var node:BinaryRedBlackTree? = tree
        var parent:BinaryRedBlackTree?
        var result:Int = 0
        while node != nil{
            parent = node
            result = comparator(value,node!.value)
            if result < 0 {
                node = node!.left
            }else if result > 0 {
                node = node!.right
            }else{
                
            }
        }
        let newNode = BinaryRedBlackTree.init(num: value, parentNode: parent)
        if result < 0 {
            parent?.left = newNode
        }else if result > 0{
            parent?.right = newNode
        }
        afterAdd(node: newNode)
    }
    /// 根据元素获取节点
    func node(withElement el:Int) -> BinaryRedBlackTree? {
        var node = root
        while node != nil {
            let result = el - node!.value
            if result == 0 {
                return node
            }else if result > 0{
                node = node!.right
            }else{
                node = node!.left
            }
        }
        return nil
    }
    /// 移除
    public func remove(element:Int){
        guard let node = node(withElement: element) else {
            return
        }
        removeNode(node: node)
    }
    func removeNode(node:BinaryRedBlackTree) {
        var tree = node
        if tree.hasTwoChild() { // 度为2的节点
            // 找到后继节点
            let s = successor(node: tree)
            tree.value = s.value
            tree = s
        }
        // 删除tree （tree的度必然为1或0）
        let replace = tree.left != nil ? tree.left : tree.right
        if replace != nil { // 度为1的节点
            // 更该parent
            replace?.parent = tree.parent
            if tree.parent == nil { // tree是度为1的节点并且是根节点
                root = replace
            }else if(tree.value == tree.parent?.left?.value){
                tree.parent?.left = replace
            }else{
                tree.parent?.right = replace
            }
            // 删除节点之后处理
            afterRemove(tree,replace)
        }else if (tree.parent == nil){ // tree是叶子节点并且是根节点 且度为0
            root = nil
            // 删除节点之后处理
            afterRemove(tree,nil)
        }else{ // tree 是叶子节点但不是根节点  且度为0
            if tree.value == tree.parent!.right?.value {
                tree.parent?.right =  nil
            }else{
                tree.parent?.left = nil
            }
            // 删除节点之后处理
            afterRemove(tree,nil)
        }
        
    }
    
    func afterRemove(_ node:BinaryRedBlackTree,_ replace:BinaryRedBlackTree?) {
        /// 红色节点直接返回
        if node.isRed(){return}
        /// 用以取代node的节点是红色
        if replace?.isRed() ?? false {
            replace?.blackNode()
            return
        }
        /// 删除的是根节点
        guard let parent = node.parent else {
            return
        }
        ///删除的黑色叶子节点
        /// 判断被删除的node是左还是右
        let parentLeft = parent.left == nil || node.isLeft()
        var  sibling = parentLeft ? parent.right : parent.left
        if parentLeft {
            if sibling?.isRed() ?? false {/// 兄弟节点是红色
                sibling?.blackNode()
                parent.redNode()
                rotateRight(node: parent)
                sibling = parent.right
            }
            /// 兄弟节点是黑色
            if sibling?.right?.isBlack() ?? true && sibling?.left?.isBlack() ?? true{
                // 兄弟节点没有一个红色子节点
                let parentBlack = parent.isBlack()
                parent.blackNode()
                sibling?.redNode()
                if parentBlack {
                    afterRemove(parent,  nil)
                }
            }else{
                /// 兄弟节点至少有一个红色字节点
                if sibling?.right?.isBlack() ?? true {
                    rotateRight(node: sibling!)
                    sibling = parent.right
                }
                sibling?.color = parent.color
                sibling?.right?.blackNode()
                parent.blackNode()
                rotateLeft(node: parent)
            }
        }else{
            if sibling?.isRed() ?? false {/// 兄弟节点是红色
                sibling?.blackNode()
                parent.redNode()
                rotateRight(node: parent)
                sibling = parent.left
            }
            /// 兄弟节点是黑色
            if sibling?.left?.isBlack() ?? true && sibling?.right?.isBlack() ?? true{
                let parentBlack = parent.isBlack()
                parent.blackNode()
                sibling?.redNode()
                if parentBlack {
                    afterRemove(parent,  nil)
                }
            }else{
                /// 兄弟节点至少有一个红色字节点
                if sibling?.left?.isBlack() ?? true {
                    rotateLeft(node: sibling!)
                    sibling = parent.left
                }
                sibling?.color = parent.color
                sibling?.left?.blackNode()
                parent.blackNode()
                rotateRight(node: parent)
            }
        }
    }
    /// 找前驱节点
    func predecssor(node:BinaryRedBlackTree) -> BinaryRedBlackTree{
        var tree = node.left
        var result = tree
        while tree != nil {
            result = tree
            tree = tree?.right
        }
        return result!
    }
    /// 找后继节点
    func successor(node:BinaryRedBlackTree) -> BinaryRedBlackTree{
        var tree = node.right
        var result = tree
        while tree != nil {
            result = tree
            tree = tree?.left
        }
        return result!
    }
    
    /// 添加之后处理
    func afterAdd(node:BinaryRedBlackTree){
        /// 添加是根节点
        guard let parent = node.parent else {
            node.blackNode()
            return
        }
        /// 如果父节点是黑色直接返回
        if parent.isBlack() {return}
        guard let grand = parent.parent else {return}
        
        /// 父节点是红色   存在叔父节点
        ///  uncle 节点 / 祖父节点
        if let uncle = parent.sibling(){
            if uncle.isRed(){ // 叔父节点是红色
                uncle.blackNode()
                parent.blackNode()
                grand.redNode()
                afterAdd(node: grand)
                return
            }
        }
        /// （B树种的叔父节点是黑色，二叉搜索树种的grand节点） 不存在叔父节点
        if parent.isLeft() {// L
            if node.isLeft() { // LL
                grand.redNode()
                parent.blackNode()
                rotateRight(node: grand)
            }else{// LR
                node.blackNode()
                grand.redNode()
                rotateLeft(node: parent)
                rotateRight(node: grand)
            }
        }else{// R
            if node.isLeft() {// RL
                node.blackNode()
                grand.redNode()
                rotateRight(node: parent)
                rotateLeft(node: grand)
            }else{// RR
                grand.redNode()
                parent.blackNode()
                rotateLeft(node: node)
            }
        }
    }
    /// 右旋
    func rotateRight(node:BinaryRedBlackTree){
        let grand = node
        let parent = grand.left
        
        let child = parent?.right
        
        grand.left = child
        parent?.right = grand
        parent?.parent = grand.parent
        
        if grand.isLeft() {
            grand.parent?.left = parent
        }else if grand.isRight(){
            grand.parent?.right = parent
        }else{
            root = parent
        }
        grand.parent = parent
        child?.parent = grand
    }
    /// 左旋
    func rotateLeft(node:BinaryRedBlackTree){
        let grand = node
        let parent = grand.right
        let child = parent?.left
        
        grand.right = child
        parent?.left = grand
        parent?.parent = grand.parent

        if grand.isLeft() {
            grand.parent?.left = parent
        }else if grand.isRight(){
            grand.parent?.right = parent
        }else{
            root = parent
        }
        
        grand.parent = parent
        child?.parent = grand
        
    }
}

class BinaryRedBlackTree:Equatable {
    static func == (lhs: BinaryRedBlackTree, rhs: BinaryRedBlackTree) -> Bool {
        return lhs.value == rhs.value
    }
    /// color
    var color:String = "red"
    var value:Int
    var parent:BinaryRedBlackTree?
    var left:BinaryRedBlackTree?
    var right:BinaryRedBlackTree?
    
    /// 初始化
    init(num:Int,parentNode:BinaryRedBlackTree?) {
        value = num
        parent = parentNode
    }
    /// 是否是左子树
    func isLeft() -> Bool {
        if value == parent?.left?.value {
            return true
        }
        return false
    }
    /// 是否是右子树
    func isRight() -> Bool {
        if value == parent?.right?.value {
            return true
        }
        return false
    }
    /// 兄弟节点
    func sibling() -> BinaryRedBlackTree?{
        if isLeft() {
            return parent?.right
        }
        if isRight() {
            return parent?.left
        }
        return nil
    }
    /// 染色  red
    func redNode(){
        color = "red"
    }
    /// 染色  black
    func blackNode(){
        color = "black"
    }
    /// 是否为红色
    func isRed() -> Bool {
        if color == "red" {
            return true
        }
        return false
    }
    /// 是否为黑色
    func isBlack() -> Bool {
        if color == "black" {
            return true
        }
        return false
    }
    func isLeftChild() -> Bool {
        return parent != nil && self == parent?.left
    }
    
    func isRightChild() -> Bool {
        return parent != nil && self == parent?.right
    }
    func hasTwoChild() -> Bool {
        return left != nil && right != nil
    }
    
    /// 打印函数
    var asString:String { return treeString(self){("\($0.color)_\($0.value)",$0.left,$0.right)}  }

}
